#!/bin/bash
#
# Copyright (c) Citrix Systems 2008. All rights reserved.
#
# Script which lets you edit the grub.conf in a Linux guest virtual disk.

if [ ! -e @INVENTORY@ ]; then
  echo Must run on a Citrix Hypervisor host.
  exit 1
fi

. @INVENTORY@

XE="@BINDIR@/xe"

device_number=
grub_conf=
default_file_list="/boot/grub/grub.cfg /grub/grub.cfg /boot/grub2/grub.cfg /grub2/grub.cfg"
default_file_list="$default_file_list /boot/isolinux/isolinux.cfg /boot/extlinux/extlinux.conf /boot/extlinux.conf /extlinux/extlinux.conf /extlinux.conf"
default_file_list="$default_file_list /boot/grub/menu.lst /boot/grub/grub.conf /grub/menu.lst"

function usage {
    echo "Usage: $0 [-u <VM UUID>] [-n <VM name>]"
    echo "    [-p <partition number>] [-f <filename>]"
    echo
    echo " -u: UUID of the VM boot disk you wish to edit"
    echo " -n: Name label of the VM boot disk you wish to edit"
    echo " -p: Partition number to mount (default: none)"
    echo " -f: Location of bootloader config file to edit"
    echo "       (defaults: ${default_file_list})"
    echo
    echo "Either -u or -n must be supplied to specify the VM to edit"
    echo "Set the EDITOR environment variable if desired (default: nano)"
    exit 1
}

while getopts "hu:n:p:f:" opt ; do
    case $opt in
    h) usage ;;
    u) vm_uuid=${OPTARG} ;;
    n) vm_name="${OPTARG}" ;;
    p) device_number=p${OPTARG} ;;
    f) default_file_list="${OPTARG}" ;;
    *) echo "Invalid option"; usage ;;
    esac
done

# determine a VM UUID
IFS=,
if [ "${vm_name}" != "" ]; then
   vm_uuids=$(${XE} vm-list name-label="${vm_name}" params=uuid --minimal)
   num_vm_uuids=0
   for vm in ${vm_uuids}; do
      num_vm_uuids=$(( ${num_vm_uuids} + 1 ))
   done
   if [ ${num_vm_uuids} -eq 0 ]; then
      echo Unable to find VMs matching name-label: ${vm_name}
      exit 1
   elif [ ${num_vm_uuids} -gt 1 ]; then
      echo Found multiple VMs with name-label: ${vm_name}
      echo ${vm_uuids}
      echo Please specify one of the UUIDs using the -u option
      exit 1
   fi
   vm_uuid=${vm_uuids}
fi

# verify that the VM UUID is valid
if [ "${vm_uuid}" = "" ]; then usage; fi

checkuuid=$(${XE} vm-list uuid="${vm_uuid}" params=uuid --minimal)
if [ "${checkuuid}" = "" ]; then
    echo Specified VM UUID "${vm_uuid}" does not exist.
    exit 1
fi

# select nano as the default editor
if [ "${EDITOR}" = "" ]; then
   EDITOR=nano
fi

mnt=
function cleanup {
   if [ ! -z "${mnt}" ]; then
      umount ${mnt} >/dev/null 2>&1
      rmdir ${mnt}
   fi

   kpartx -dvspp ${device}

   if [ ! -z "${vbd_uuid}" ]; then
      echo -n "Unplugging VBD: "
      ${XE} vbd-unplug uuid=${vbd_uuid} timeout=20
      # poll for the device to go away if we know its name
      if [ "${device}" != "" ]; then
          device_gone=0
          for ((i=0; i<10; i++)); do
              echo -n "."
              if [ ! -b ${device} ]; then
                  echo " done"
                  device_gone=1
                  break
              fi
              sleep 1
          done
          if [ ${device_gone} -eq 0 ]; then
              echo " failed"
              echo Please destroy VBD ${vbd_uuid} manually.
          else
              ${XE} vbd-destroy uuid=${vbd_uuid}
          fi
      fi
   fi
}

# Find the bootable VBD for the VM
IFS=","
bootable_vbds=$(${XE} vbd-list vm-uuid=${vm_uuid} bootable=true params=uuid --minimal)
num_bootable_vbds=0
for vbd in ${bootable_vbds}; do
   num_bootable_vbds=$(( ${num_bootable_vbds} + 1 ))
done
if [ ${num_bootable_vbds} -eq 0 ]; then
   echo Unable to find a bootable disk for VM.
   echo There are no VBDs marked with bootable=true.
   exit 1
elif [ ${num_bootable_vbds} -gt 1 ]; then
   echo Found ${num_bootable_vbds} for VM, but only 1 is allowed.
   echo Please remove the bootable=true flag from some VBDs to fix this.
   exit 1
fi

# we now know that there is only one bootable VBD, so use that
vbd_uuid=${bootable_vbds}
vdi_uuid=$(${XE} vbd-param-get uuid=${vbd_uuid} param-name=vdi-uuid)
if [ "${vdi_uuid}" = "" ]; then
    echo Unable to determine VDI associated with VBD ${vbd_uuid}
    exit 1
fi

# check that the VDI is not currently in use
vbd_uuids=$(${XE} vbd-list vdi-uuid=${vdi_uuid} params=uuid --minimal)
if [ $? -ne 0 ]; then
   echo Error in vbd-list call
   exit 1
fi

for vbd in ${vbd_uuids}; do
   attached=$(${XE} vbd-list params=currently-attached uuid=${vbd} --minimal)
   case ${attached} in
   true)
     vm_uuid=$(${XE} vbd-list params=vm-uuid uuid=${vbd} --minimal)
     echo The VDI is currently in use by VBD:
     echo ${vbd}
     echo The associated VM is:
     echo ${vm_uuid}
     echo
     ${XE} vm-list uuid=${vm_uuid}
     echo Please shut down any running VMs which are currently using the disk.
     exit 1
     ;;
   false)
     ;;
   *)
     echo Internal error: unknown currently-attached result: ${attached}
     ;;
   esac
done

echo -n "Creating dom0 VBD: "
vbd_uuid=$(${XE} vbd-create vm-uuid=${CONTROL_DOMAIN_UUID} vdi-uuid=${vdi_uuid} device=autodetect)
echo ${vbd_uuid}

if [ $? -ne 0 -o -z "${vbd_uuid}" ]; then
  echo error creating VBD for dom0 block attach
  cleanup
  exit 1
fi

echo -n "Plugging VBD: "
${XE} vbd-plug uuid=${vbd_uuid}

device=/dev/$(${XE} vbd-param-get uuid=${vbd_uuid} param-name=device)
echo ${device}


if [ ! -b ${device} ]; then
  echo ${device}: not a block special
  cleanup
  exit 1
fi

kpartx -avspp ${device}
mapped_device="/dev/mapper/${vdi_uuid}"

echo -n "Waiting for ${mapped_device}${device_number}: "
found_device=0

for ((i=0; i<5; i++)); do
   echo -n '.'
   if [ -b ${mapped_device}${device_number} ]; then
      found_device=1
      echo ' done'
      break
   fi
   sleep 1
done

if [ ${found_device} -eq 0 ]; then
   echo Device ${mapped_device}${device_number} not found.
   echo You must specify the correct partition number with -p
   cleanup
   exit 1
fi

echo -n "Mounting filesystem: "
mnt=/var/run/fix-grub-${vdi_uuid}
mkdir -p ${mnt}

mount ${mapped_device}${device_number} ${mnt} > /dev/null 2>&1

if [ $? -ne 0 ]; then
  echo " failed"
  echo Partitions in the VDI are:
  echo
  ls -la ${mapped_device}*
  echo
  echo You can use the -p option to specify a partition number to mount.
  cleanup
  exit 1
fi

echo " done"

IFS=" "
if [ "${grub_conf}" = "" ]; then
   for f in ${default_file_list}; do
      if [ -e "${mnt}${f}" ]; then
          grub_conf="${f}"
          break;
      fi
   done
fi

if [ "${grub_conf}" = "" ]; then
   echo Unable to find a valid bootloader config file in:
   echo "  ${default_file_list}"
   echo Please specify one with the -f option
   cleanup
   exit 1
fi

${EDITOR} ${mnt}${grub_conf}

cleanup
